import torch
import numpy as np
import cv2 as cv


def letter_box(img, new_shape):
    """将图片等比例缩放调整到指定边长的正方形,剩下的填充"""
    shape = img.shape[:2]  # [h, w]
    r = min(new_shape / shape[0], new_shape / shape[1])  # scale ratio (new / old)
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    img = cv.resize(img, new_unpad, interpolation=cv.INTER_LINEAR)
    dw, dh = (new_shape - new_unpad[0]) / 2, (new_shape - new_unpad[1]) / 2  # wh padding
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))  # 计算上下两侧的padding
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))  # 计算左右两侧的padding
    img = cv.copyMakeBorder(img, top, bottom, left, right, cv.BORDER_CONSTANT, value=(0, 0, 0))  # add border
    return img


def image_pretreat(img, shape):
    img = letter_box(img, shape)
    img = img[:, :, ::-1].transpose(2, 0, 1)
    return img / 256


def image_translation(img, dx, dy):
    affine_arr = np.float32([[1, 0, dx], [0, 1, dy]])
    return cv.warpAffine(img, affine_arr, (img.shape[0], img.shape[1]))


def image_show(img, pred):
    img = img.numpy().transpose(1, 2, 0)
    img = img[:, :, ::-1] * 256
    img = np.ascontiguousarray(img).astype('uint8')
    img_shape = 224
    conf = torch.sigmoid(pred[0]).numpy()
    if conf > 0.4:
        x = (torch.sigmoid(pred[1]) * img_shape).numpy()
        y = (torch.sigmoid(pred[2]) * img_shape).numpy()
        w = (0.4 * torch.exp(pred[3]) * img_shape).numpy()
        h = (0.6 * torch.exp(pred[4]) * img_shape).numpy()
        x1 = np.floor((x - w / 2) + 0.5).astype('int32')
        y1 = np.floor((y - h / 2) + 0.5).astype('int32')
        x2 = np.floor((x + w / 2) + 0.5).astype('int32')
        y2 = np.floor((y + h / 2) + 0.5).astype('int32')
        cv.rectangle(img, (x1, y1), (x2, y2), color=(0, 0, 255), thickness=2)
        cv.putText(img, '{:.4f}'.format(conf), (x1, y1), cv.FONT_HERSHEY_PLAIN, 1.2, (0, 0, 255), 1)
    cv.imshow('test', img)
    cv.waitKey(10)


def freeze_param(model, exclude='none'):
    """冻结网络参数
    :param model:网络模型
    :param exclude:哪部分不冻结,'classify'或'bbr'
    """
    if exclude == 'classify':
        model.conv.weight.requires_grad = False
        model.conv.bias.requires_grad = False
        model.fc2.weight.requires_grad = False
        model.fc2.bias.requires_grad = False
    elif exclude == 'bbr':
        for param in model.parameters():
            param.requires_grad = False
        for param in model.backbone.parameters():
            param.requires_grad = False
        model.conv.weight.requires_grad = True
        model.conv.bias.requires_grad = True
        model.fc2.weight.requires_grad = True
        model.fc2.bias.requires_grad = True
    else:
        print('Wrong exclude param')
